"""
异步 Python 客户端示例
使用 aiohttp 进行异步 API 调用，适合高并发场景
"""
import aiohttp
import asyncio
from typing import Optional, Dict, List


class AsyncFastAPIClient:
    """异步 FastAPI 客户端类"""
    
    def __init__(self, base_url: str = "http://localhost:8000/api"):
        """
        初始化异步客户端
        
        Args:
            base_url: API 基础地址
        """
        self.base_url = base_url
        self.token: Optional[str] = None
        self.session: Optional[aiohttp.ClientSession] = None
    
    async def __aenter__(self):
        """异步上下文管理器入口"""
        self.session = aiohttp.ClientSession()
        return self
    
    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """异步上下文管理器出口"""
        if self.session:
            await self.session.close()
    
    def _get_headers(self, include_auth: bool = True) -> Dict[str, str]:
        """获取请求头"""
        headers = {"Content-Type": "application/json"}
        if include_auth and self.token:
            headers["Authorization"] = f"Bearer {self.token}"
        return headers
    
    async def register(self, username: str, email: str, password: str) -> Dict:
        """用户注册"""
        url = f"{self.base_url}/auth/register"
        data = {
            "username": username,
            "email": email,
            "password": password
        }
        
        try:
            async with self.session.post(
                url,
                json=data,
                headers=self._get_headers(include_auth=False)
            ) as response:
                if response.status == 201:
                    result = await response.json()
                    return {"success": True, "data": result}
                else:
                    error = await response.json()
                    return {"success": False, "error": error.get("detail", "注册失败")}
        except Exception as e:
            return {"success": False, "error": str(e)}
    
    async def login(self, username: str, password: str) -> Dict:
        """用户登录"""
        url = f"{self.base_url}/auth/login"
        data = aiohttp.FormData()
        data.add_field("username", username)
        data.add_field("password", password)
        
        try:
            async with self.session.post(url, data=data) as response:
                if response.status == 200:
                    result = await response.json()
                    self.token = result.get("access_token")
                    return {"success": True, "data": result}
                else:
                    error = await response.json()
                    return {"success": False, "error": error.get("detail", "登录失败")}
        except Exception as e:
            return {"success": False, "error": str(e)}
    
    async def get_current_user(self) -> Dict:
        """获取当前用户信息"""
        url = f"{self.base_url}/auth/me"
        
        try:
            async with self.session.get(
                url,
                headers=self._get_headers()
            ) as response:
                if response.status == 200:
                    result = await response.json()
                    return {"success": True, "data": result}
                else:
                    error = await response.json()
                    if response.status == 401:
                        self.token = None
                    return {"success": False, "error": error.get("detail", "获取失败")}
        except Exception as e:
            return {"success": False, "error": str(e)}
    
    async def get_users(self, skip: int = 0, limit: int = 100) -> Dict:
        """获取用户列表"""
        url = f"{self.base_url}/users/"
        params = {"skip": skip, "limit": limit}
        
        try:
            async with self.session.get(
                url,
                params=params,
                headers=self._get_headers(include_auth=False)
            ) as response:
                if response.status == 200:
                    result = await response.json()
                    return {"success": True, "data": result}
                else:
                    return {"success": False, "error": "获取失败"}
        except Exception as e:
            return {"success": False, "error": str(e)}


async def example_async_basic():
    """异步基础示例"""
    print("=" * 50)
    print("异步客户端基础示例")
    print("=" * 50)
    
    async with AsyncFastAPIClient() as client:
        # 注册
        print("1. 注册用户...")
        result = await client.register(
            username="asyncuser",
            email="async@example.com",
            password="async123"
        )
        if result["success"]:
            print(f"   ✅ 注册成功：{result['data']['username']}")
        else:
            print(f"   ⚠️ {result['error']}")
        
        # 登录
        print("\n2. 登录...")
        result = await client.login("asyncuser", "async123")
        if result["success"]:
            print(f"   ✅ 登录成功")
            print(f"   Token: {result['data']['access_token'][:30]}...")
        
        # 获取用户信息
        print("\n3. 获取用户信息...")
        result = await client.get_current_user()
        if result["success"]:
            user = result["data"]
            print(f"   ✅ 用户：{user['username']} ({user['email']})")


async def example_async_concurrent():
    """异步并发示例：同时获取多个用户"""
    print("\n" + "=" * 50)
    print("异步并发示例：批量获取用户")
    print("=" * 50)
    
    async with AsyncFastAPIClient() as client:
        # 先获取用户列表
        result = await client.get_users(limit=10)
        if result["success"]:
            users = result["data"]
            print(f"获取到 {len(users)} 个用户，并发获取详细信息...\n")
            
            # 并发获取每个用户的详细信息（使用异步方法）
            tasks = []
            for user in users[:5]:  # 只获取前5个
                # 这里可以直接使用同步请求在异步环境中
                url = f"{client.base_url}/users/{user['id']}"
                task = client.session.get(url, headers=client._get_headers(include_auth=False))
                tasks.append(task)
            
            # 等待所有请求完成
            responses = await asyncio.gather(*tasks, return_exceptions=True)
            
            print("并发获取结果：")
            for i, response in enumerate(responses):
                if isinstance(response, Exception):
                    print(f"  {i+1}. 请求异常: {response}")
                else:
                    async with response:
                        if response.status == 200:
                            user = await response.json()
                            print(f"  {i+1}. {user['username']} ({user['email']})")
                        else:
                            print(f"  {i+1}. 获取失败 (状态码: {response.status})")


async def main():
    """主函数"""
    await example_async_basic()
    await example_async_concurrent()


if __name__ == "__main__":
    print("\n异步 FastAPI 客户端示例\n")
    asyncio.run(main())

